#!/usr/bin/python3
# coding=utf-8

import numpy as np
import os

def flip_h(x):
    """沿H维度翻转张量（上下翻转）
    
    参数:
        x: 输入张量，形状为 [N, C, H, W]
        
    返回:
        沿H维度翻转后的张量
    """
    # 沿第2个轴（H维度）翻转
    return np.flip(x, axis=3)

def gen_golden_data_flip_h():
    """生成H维度Flip操作的测试数据"""
    # 配置参数
    dtype = np.float32
    input_shape = [9, 10, 23, 1971]
    
    # 生成输入数据 - 使用有规律的数据便于验证
    np.random.seed(42)  # 固定随机种子以便复现
    x = np.arange(np.prod(input_shape)).reshape(input_shape).astype(dtype)
    
    # 创建输入输出目录
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # 计算H维度翻转结果
    golden = flip_h(x)
    
    # 保存二进制文件
    x.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")
    
    # 保存文本文件
    with open("./input/input_x.txt", "w") as f:
        f.write("Input tensor shape: {}\n".format(input_shape))
        f.write("Input data:\n")
        # 将数据展平后写入，每行8个元素
        flat_x = x.flatten()
        for i in range(0, len(flat_x), 8):
            line = " ".join("{:8.2f}".format(val) for val in flat_x[i:i+8])
            f.write(line + "\n")
    
    with open("./output/golden.txt", "w") as f:
        f.write("Output tensor shape: {}\n".format(golden.shape))
        f.write("Output data (flipped along H dimension):\n")
        # 将数据展平后写入，每行8个元素
        flat_golden = golden.flatten()
        for i in range(0, len(flat_golden), 8):
            line = " ".join("{:8.2f}".format(val) for val in flat_golden[i:i+8])
            f.write(line + "\n")
    
    # 同时保存一个详细的对比文件
    with open("./output/comparison.txt", "w") as f:
        f.write("Comparison between input and output:\n")
        f.write("=" * 60 + "\n")
        
        # 由于是H维度翻转，我们按H维度来显示对比
        f.write("Original data (H dimension):\n")
        for h in range(x.shape[2]):  # H维度
            f.write("H={}: ".format(h))
            line = " ".join("{:6.1f}".format(val) for val in x[0, 0, h, :8])  # 只显示前8个W维度
            f.write(line)
            if len(x[0, 0, h]) > 8:
                f.write(" ...")
            f.write("\n")
        
        f.write("\nFlipped data (H dimension):\n")
        for h in range(golden.shape[2]):  # H维度
            f.write("H={}: ".format(h))
            line = " ".join("{:6.1f}".format(val) for val in golden[0, 0, h, :8])  # 只显示前8个W维度
            f.write(line)
            if len(golden[0, 0, h]) > 8:
                f.write(" ...")
            f.write("\n")

if __name__ == "__main__":
    gen_golden_data_flip_h()